from openai import OpenAI
import qdrant_client
from IPython.display import Markdown, display
import os
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if project_root not in sys.path:
    sys.path.append(project_root)
from rag.config.config import Config
from rag.generator.utils import get_generator
from rag.reranker.utils import get_reranker
from rag.retriever.utils import get_retriever
from pymilvus import model

def query_qdrant(query, collection_name, vector_name="record", top_k=5):
    api_key="sk-feebadbcb5654f5c8f9044c78f7c4548"
    client = OpenAI(api_key=api_key,
                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
    EMBEDDING_MODEL = "text-embedding-v2"
    # Creates embedding vector from user query
    embedded_query = (
        client.embeddings.create(
            input=query,
            model=EMBEDDING_MODEL,
        )
       .data[0]
       .embedding
    )
    qdrant = qdrant_client.QdrantClient(host="localhost", port=6333)
    query_results = qdrant.search(
        collection_name=collection_name,
        query_vector=(vector_name, embedded_query),
        limit=top_k,
    )
    return query_results


def query_docs(query):
    """Query the knowledge base for relevant articles."""
    print(f"Searching knowledge base with query: {query}")
    query_results = query_qdrant(query, collection_name="Customs declaration records")
    output = []
    for i, article in enumerate(query_results):
        declaration_number = article.payload["报关单号"]
        shipper = article.payload["发货人"]
        departure_place = article.payload["发货地"]
        destination_place = article.payload["收货地"]
        goods_name = article.payload["货物名称"]
        transportation_mode = article.payload["运输方式"]
        declaration_date = article.payload["报关日期"]

        output.append((declaration_number, shipper, departure_place, destination_place, goods_name, transportation_mode, declaration_date))

    if output:
        declaration_number, shipper, departure_place, destination_place, goods_name, transportation_mode, declaration_date = output[0]
        response = f"报关单号: {declaration_number}\n发货人: {shipper}\n发货地: {departure_place}\n收货地: {destination_place}\n货物名称: {goods_name}\n运输方式: {transportation_mode}\n报关日期: {declaration_date}"
        print("Most relevant record:")
        print(response)
        return {"response": response}
    else:
        print("No results")
        return {"response": "No results found."}

def milvus_search(query):
    embedding_model = model.hybrid.BGEM3EmbeddingFunction(
        devices="cuda:0", return_sparse=True, return_dense=True, return_colbert_vecs=False
    )
    query = [query]

    config = Config()

    retriever = get_retriever(config)
    retrieved_list = retriever.retrieve(query)

    reranker = get_reranker(config)
    reranked_list = reranker.rerank(query, retrieved_list)

    generator = get_generator(config)
    answer = generator.generate(query, reranked_list)
    return answer